import os
import cv2
import glob
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import ocr.segment_model as osm

def resize_image(im):
    h, w, _ = im.shape
    size = (int(512), int(512))
    im = cv2.resize(im, size, interpolation=cv2.INTER_AREA)
    # la_p = cv2.resize(label_im, size, interpolation=cv2.INTER_AREA)

    ratio_h = 512 / float(h)
    ratio_w = 512 / float(w)

    return im, (ratio_h, ratio_w)

class Detector(object):
    def __init__(self, model_dir):
        os.environ['CUDA_VISIBLE_DEVICES'] = '1'
        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.per_process_gpu_memory_fraction = 1.0
        config.gpu_options.allow_growth = True
        self.input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images')
        self.session = tf.Session(config=config)
        with tf.variable_scope("", reuse=tf.AUTO_REUSE):
            self.global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)
            self.score_nrow, self.score_ncol, self.score_row, self.score_col = osm.model(self.input_images, is_training=True)
            self.variable_averages = tf.train.ExponentialMovingAverage(0.997, self.global_step)
            self.saver = tf.train.Saver(self.variable_averages.variables_to_restore())
            self.ckpt_state = tf.train.get_checkpoint_state(model_dir)
            self.model_path = os.path.join(model_dir, os.path.basename(self.ckpt_state.model_checkpoint_path))
            self.saver.restore(self.session, self.model_path)

    def main_detection(self, image):
        im_resized, (ratio_h, ratio_w) = resize_image(image)
        # print('im_resize shape:', im_resized.shape)
        # cv2.imwrite('resize.jpg', im_resized)
        # im_resized = cv2.resize(im_resized, dsize=(512, 512), interpolation=cv2.INTER_AREA)
        # 将图像由BGR转化为RGB模式输入
        im_resized = im_resized[:, :, ::-1].astype(np.float32)
        score_nrow, score_ncol, score_row, score_col = self.session.run([self.score_nrow, self.score_ncol, self.score_row, self.score_col], feed_dict={self.input_images: [im_resized]})
        # plt.subplot(2, 2)
        # plt.imshow(0, 0, score_nrow[0,:,:,0])
        # plt.imshow(0, 1, score_ncol[0,:,:,0])
        # plt.imshow(1, 0, score_row[0,:,:,0])
        # plt.imshow(1, 1, score_col[0,:,:,0])
        # plt.show()
        return score_nrow[0], score_ncol[0], score_row[0], score_col[0], ratio_h, ratio_w

def generate_map(image_path, model_path):

    instance = Detector(model_path)
    # images = glob.glob(image_path)

    image_color = cv2.resize(image_path, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    score_nrow, score_ncol, score_row, score_col, ratio_h, ratio_w = instance.main_detection(image_color)

    score_nrow = np.where(score_nrow > 0.9, score_nrow, 0)
    score_nrow = np.where(score_nrow < 0.9, score_nrow, 1)

    score_ncol = np.where(score_ncol > 0.9, score_ncol, 0)
    score_ncol = np.where(score_ncol < 0.9, score_ncol, 1)

    score_row = np.where(score_row > 0.9, score_row, 0)
    score_row = np.where(score_row < 0.9, score_row, 1)

    score_col = np.where(score_col > 0.9, score_col, 0)
    score_col = np.where(score_col < 0.9, score_col, 1)

    nmap = cv2.bitwise_and(score_nrow, score_ncol)
    lmap = cv2.bitwise_and(score_row, score_col)

    score_nrow_map = cv2.resize(score_nrow, dsize=None, fx=1/ratio_w, fy=1/ratio_h, interpolation=cv2.INTER_AREA)
    score_ncol_map = cv2.resize(score_ncol, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
    score_row_map = cv2.resize(score_row, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
    score_col_map = cv2.resize(score_col, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)

    return score_nrow_map, score_ncol_map, score_row_map, score_col_map

def pbInference(image_path, graph_path):
    graph_def = tf.GraphDef()
    with tf.gfile.FastGFile(graph_path, 'rb') as f:
        graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')

    # tensor_name_list = [tensor.name for tensor in tf.get_default_graph().as_graph_def().node]

    isess = tf.InteractiveSession()
    images_palceholder = tf.get_default_graph().get_tensor_by_name("input_images:0")
    feat1 = tf.get_default_graph().get_tensor_by_name("model_0/feature_fusion/Conv_7/Sigmoid:0")
    feat2 = tf.get_default_graph().get_tensor_by_name("model_0/feature_fusion/Conv_8/Sigmoid:0")
    feat3 = tf.get_default_graph().get_tensor_by_name("model_0/feature_fusion/Conv_9/Sigmoid:0")
    feat4 = tf.get_default_graph().get_tensor_by_name("model_0/feature_fusion/Conv_10/Sigmoid:0")

    image = cv2.resize(image_path, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    ratio_w, ratio_h = 1.0, 1.0
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image[:, :, ::-1].astype(np.float32)
    image = image[np.newaxis, :]
    score_nrow, score_ncol, score_row, score_col = isess.run([feat1, feat2, feat3, feat4], feed_dict={images_palceholder: image})   # 耗时近5.4秒
    # cv2.namedWindow('score_map', cv2.WINDOW_NORMAL)
    # cv2.imshow('score_map', score_row[0, :, :, 0]*255)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    score_nrow, score_ncol, score_row, score_col = score_nrow[0], score_ncol[0], score_row[0], score_col[0]

    score_nrow = np.where(score_nrow > 0.9, score_nrow, 0)
    score_nrow = np.where(score_nrow < 0.9, score_nrow, 1)

    score_ncol = np.where(score_ncol > 0.9, score_ncol, 0)
    score_ncol = np.where(score_ncol < 0.9, score_ncol, 1)

    score_row = np.where(score_row > 0.9, score_row, 0)
    score_row = np.where(score_row < 0.9, score_row, 1)

    score_col = np.where(score_col > 0.9, score_col, 0)
    score_col = np.where(score_col < 0.9, score_col, 1)

    nmap = cv2.bitwise_and(score_nrow, score_ncol)
    lmap = cv2.bitwise_and(score_row, score_col)

    score_nrow_map = cv2.resize(score_nrow, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
    score_ncol_map = cv2.resize(score_ncol, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
    score_row_map = cv2.resize(score_row, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
    score_col_map = cv2.resize(score_col, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)

    return score_nrow_map, score_ncol_map, score_row_map, score_col_map
